import pyscf

from pyscf import gto, scf
from pyscf.gto import Mole
from gpu4pyscf.dft import rks, uks

from typing import List, Tuple, Union, Dict, Any

from requests import get
import torch
from torch import Tensor

from . import ForceCell
from ..system import Molecule
from ..function.units import *
from ..function import gather_value, get_ms_array, get_arguments

class PySCFForceCell(ForceCell):
    def __init__(self, 
                 system: Molecule, 
                 pyscf_option: str = 'ks', # hf, ks, sqm
                 mol_parameters: Union[Dict, Mole]=None, 
                 ks_parameters: Union[Dict, str]=None,
                 length_unit: str = None, 
                 energy_unit: str = None, 
                 use_pbc: bool = None,):
        super().__init__(length_unit=length_unit, 
                         energy_unit=energy_unit, 
                         use_pbc=use_pbc)
        self.atom_type = system.atomic_number
        self.pyscf_option = pyscf_option
        if mol_parameters is Mole:
            self.pyscf_mol = mol_parameters
        elif mol_parameters is None:
            self.pyscf_mol = gto.Mole()
            self.pyscf_mol.basis = '6-31g'
            self.pyscf_mol.unit = 'a'
        elif type(mol_parameters) is dict:
            self.pyscf_mol = gto.Mole(**mol_parameters)
        self.dm0 = None
        self.coord_fac = length_convert(self.length_unit, 'a')
        if self.pyscf_option == "ks":
            if ks_parameters is None:
                self.ks_parameters = {'xc': 'LDA'}
            elif type(ks_parameters) is str:
                self.ks_parameters = {'xc': ks_parameters}
            else:
                self.ks_parameters = ks_parameters

        if self.pyscf_option == "sqm": 
            if pyscf.__version__[:3] == '1.7':
                from pyscf.semiempirical import RMINDO3
                mf = RMINDO3(self.pyscf_mol)
                g = mf.nuc_grad_method()
            else:
                raise ImportError("pyscf version should be 1.7.* when using semiempirical method")
            
        elif self.pyscf_option == "ks":
            if self.pyscf_mol.spin == 0:
                mf = rks.RKS(self.pyscf_mol, **self.ks_parameters)
            else:
                mf = uks.UKS(self.pyscf_mol, **self.ks_parameters)
            # energy = mf.kernel(dm0=self.dm0)
            # self.dm0 = mf.make_rdm1()

        elif self.pyscf_option == "hf":
            if self.pyscf_mol.spin == 0:
                mf = scf.RHF(self.pyscf_mol)
            else:
                mf = scf.UHF(self.pyscf_mol)
        g = mf.nuc_grad_method()

        self.mf = mf.as_scanner()
        self.g = g.as_scanner()

    def print_info(self):
        print(f"{'--'*40}")
        print("PySCF running")
        print(f"Basis: {self.pyscf_mol.basis}")
        if self.pyscf_option == "ks":
            print(f"KS_Param: {self.ks_parameters}")
        print(f"{'--'*40}")
    
    def forward(self, 
                  coordinate: Tensor, 
                  neighbour_index: Tensor = None, 
                  neighbour_mask: Tensor = None, 
                  neighbour_vector: Tensor = None, 
                  neighbour_distance: Tensor = None, 
                  pbc_box: Tensor = None):
        b,n,d = coordinate.shape
        coord = coordinate * self.coord_fac
        atom_coord_list = [[atom, tuple(coor)] for atom, coor in zip(self.atom_type.squeeze().tolist(), coord.detach().cpu().numpy().squeeze().tolist())]
        # print(atom_coord_list)

        energy, grad = self.g(atom_coord_list)
        energy = get_ms_array(energy, dtype=torch.float32, device=coordinate.device) * energy_convert('Ha', self.energy_unit)
        grad = get_ms_array(grad, dtype=torch.float32, device=coordinate.device)
        force = -grad * energy_convert('Ha', self.energy_unit) / length_convert('Bohr', self.length_unit)
        
        # (B, A, D)
        virial = None
        # if pbc_box is not None:
        #     # (B, D)
        virial = torch.zeros((b,d),dtype=torch.float32,device=coordinate.device) # not implemented

        return energy, force, virial